import sys
import os
import numpy as np
import pandas as pd
from PyQt6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                             QHBoxLayout, QPushButton, QLabel, QFrame, QFileDialog,
                             QDialog, QComboBox, QTableWidget, QTableWidgetItem, 
                             QDialogButtonBox, QMessageBox, QHeaderView, QSpinBox,
                             QColorDialog)
from PyQt6.QtCore import Qt
from PyQt6.QtGui import QColor
import pyqtgraph as pg

# --- 1. 全局样式表 ---
STYLESHEET = """
QMainWindow { background-color: #1e1e1e; }
QWidget { font-family: 'Segoe UI', 'Microsoft YaHei', sans-serif; font-size: 14px; color: #e0e0e0; }
QDialog { background-color: #252526; }

/* 侧边栏 */
QFrame#SidePanel { background-color: #252526; border-right: 1px solid #3e3e42; }

/* 按钮 */
QPushButton { background-color: #2d2d30; border: 1px solid #454545; border-radius: 4px; padding: 6px 12px; }
QPushButton:hover { background-color: #3e3e42; border-color: #007acc; }
QPushButton:pressed { background-color: #007acc; color: white; }

/* 标题 */
QLabel#Title { font-size: 16px; font-weight: bold; color: #00bcd4; margin-bottom: 10px; }

/* 表格样式 */
QTableWidget {
    background-color: #1e1e1e;
    border: 1px solid #3e3e42;
    gridline-color: #333;
    selection-background-color: #2a2d32;
    selection-color: #fff;
}
QHeaderView::section {
    background-color: #2d2d30;
    padding: 6px;
    border: 1px solid #3e3e42;
    font-weight: bold;
}

/* 复选框美化 */
QCheckBox { spacing: 5px; }
QCheckBox::indicator { 
    width: 18px; height: 18px; 
    border: 2px solid #666; 
    border-radius: 3px; 
    background: #333; 
}
QCheckBox::indicator:hover { border-color: #007acc; }

QTableWidget::indicator {
    width: 18px; height: 18px;
    border: 2px solid #888;
    background-color: #333;
    border-radius: 3px;
}
QTableWidget::indicator:checked {
    background-color: #007acc;
    border-color: #007acc;
}

QSpinBox { background-color: #333; border: 1px solid #555; color: white; padding: 2px; }
QComboBox { background-color: #333; border: 1px solid #555; padding: 2px; }
"""

# --- 2. 颜色选择按钮 ---
class ColorButton(QPushButton):
    def __init__(self, default_hex):
        super().__init__()
        self.hex_color = default_hex
        self.setText(default_hex)
        self.update_style()
        self.clicked.connect(self.pick_color)
        self.setCursor(Qt.CursorShape.PointingHandCursor)

    def update_style(self):
        c = QColor(self.hex_color)
        text_color = "black" if (c.red()*0.299 + c.green()*0.587 + c.blue()*0.114) > 150 else "white"
        self.setStyleSheet(f"""
            QPushButton {{
                background-color: {self.hex_color};
                color: {text_color};
                border: 1px solid #555;
                border-radius: 4px;
                font-weight: bold;
            }}
            QPushButton:hover {{
                border: 2px solid white;
            }}
        """)
        self.setText(self.hex_color.upper())

    def pick_color(self):
        color = QColorDialog.getColor(QColor(self.hex_color), self, "选择线条颜色")
        if color.isValid():
            self.hex_color = color.name()
            self.update_style()

# --- 3. 对话框 (无变化，保留记忆功能) ---
class ColumnSelectorDialog(QDialog):
    def __init__(self, columns, cache_config, parent=None):
        super().__init__(parent)
        self.setWindowTitle("数据列配置")
        self.resize(650, 600)
        self.selected_x = None
        self.grouped_selection = {} 
        self.new_cache_config = {} 
        
        layout = QVBoxLayout(self)

        layout.addWidget(QLabel("第一步：选择时间轴 (X Axis):"))
        self.combo_x = QComboBox()
        self.combo_x.addItems(columns)
        layout.addWidget(self.combo_x)

        layout.addSpacing(10)
        layout.addWidget(QLabel("第二步：配置数据列:"))
        
        self.table = QTableWidget()
        self.table.setColumnCount(3)
        self.table.setHorizontalHeaderLabels(["数据列名 (勾选显示)", "显示所在行 (Row ID)", "线条颜色"])
        
        header = self.table.horizontalHeader()
        header.setSectionResizeMode(0, QHeaderView.ResizeMode.Stretch)
        header.setSectionResizeMode(1, QHeaderView.ResizeMode.Interactive)
        self.table.setColumnWidth(1, 150)
        header.setSectionResizeMode(2, QHeaderView.ResizeMode.Interactive)
        self.table.setColumnWidth(2, 120)
        self.table.verticalHeader().setVisible(False)
        
        default_colors = ['#00e5ff', '#ff4081', '#ffd740', '#aeea00', '#7c4dff', '#ff6e40', '#ffffff', '#00bcd4']
        
        self.table.setRowCount(len(columns))
        
        for i, col_name in enumerate(columns):
            if col_name in cache_config:
                cfg = cache_config[col_name]
                is_checked = cfg['checked']
                row_val = cfg['row']
                color_val = cfg['color']
            else:
                is_checked = False
                row_val = i + 1
                color_val = default_colors[i % len(default_colors)]

            item_name = QTableWidgetItem(col_name)
            item_name.setFlags(Qt.ItemFlag.ItemIsUserCheckable | Qt.ItemFlag.ItemIsEnabled)
            item_name.setCheckState(Qt.CheckState.Checked if is_checked else Qt.CheckState.Unchecked)
            self.table.setItem(i, 0, item_name)
            
            sb_row = QSpinBox()
            sb_row.setRange(1, 100)
            sb_row.setValue(row_val)
            sb_row.setAlignment(Qt.AlignmentFlag.AlignCenter)
            self.table.setCellWidget(i, 1, sb_row)
            
            btn_color = ColorButton(color_val)
            self.table.setCellWidget(i, 2, btn_color)

        layout.addWidget(self.table)

        buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
        buttons.accepted.connect(self.accept_selection)
        buttons.rejected.connect(self.reject)
        layout.addWidget(buttons)

    def accept_selection(self):
        self.selected_x = self.combo_x.currentText()
        self.grouped_selection = {}
        self.new_cache_config = {}
        
        has_selection = False
        
        for row in range(self.table.rowCount()):
            item_name = self.table.item(row, 0)
            col_name = item_name.text()
            
            is_checked = (item_name.checkState() == Qt.CheckState.Checked)
            row_id = self.table.cellWidget(row, 1).value()
            hex_color = self.table.cellWidget(row, 2).hex_color
            
            self.new_cache_config[col_name] = {
                'checked': is_checked,
                'row': row_id,
                'color': hex_color
            }
            
            if is_checked:
                has_selection = True
                if row_id not in self.grouped_selection:
                    self.grouped_selection[row_id] = []
                self.grouped_selection[row_id].append({'name': col_name, 'color': hex_color})
        
        if not has_selection:
            QMessageBox.warning(self, "提示", "请至少勾选一列数据！")
            return
            
        self.accept()

# --- 4. 主窗口 ---
class ProScopeFullGrid(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("ProWave - Full Grid Edition")
        self.resize(1300, 850)
        self.setStyleSheet(STYLESHEET)

        self.df = None 
        self.plot_items = []
        self.column_state_cache = {} 
        
        self.init_ui()

    def init_ui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QHBoxLayout(central_widget)
        main_layout.setContentsMargins(0, 0, 0, 0)

        # 侧边栏
        side_panel = QFrame()
        side_panel.setObjectName("SidePanel")
        side_panel.setFixedWidth(240)
        side_layout = QVBoxLayout(side_panel)
        side_layout.setSpacing(15)
        
        side_layout.addWidget(QLabel("数据源", objectName="Title"))
        self.btn_load = QPushButton("📂 打开 CSV 文件")
        self.btn_load.clicked.connect(self.load_csv)
        side_layout.addWidget(self.btn_load)

        self.btn_select_cols = QPushButton("⚙️ 配置数据列")
        self.btn_select_cols.clicked.connect(self.open_column_selector)
        self.btn_select_cols.setEnabled(False)
        side_layout.addWidget(self.btn_select_cols)
        
        side_layout.addWidget(QLabel("视图控制", objectName="Title"))
        self.btn_reset = QPushButton("🔍 复位视图")
        self.btn_reset.clicked.connect(self.reset_views)
        side_layout.addWidget(self.btn_reset)

        side_layout.addStretch()
        
        group_test = QFrame()
        group_test.setStyleSheet("background: #2d2d30; border-radius: 5px; padding: 5px;")
        layout_test = QVBoxLayout(group_test)
        layout_test.addWidget(QLabel("测试工具:", styleSheet="color:gray; font-size:12px; border:none;"))
        self.btn_gen_test = QPushButton("⚡ 生成测试数据")
        self.btn_gen_test.setStyleSheet("border: 1px dashed #666;")
        self.btn_gen_test.clicked.connect(self.generate_test_csv)
        layout_test.addWidget(self.btn_gen_test)
        side_layout.addWidget(group_test)
        
        self.lbl_status = QLabel("就绪")
        self.lbl_status.setStyleSheet("color: #888; margin-top: 5px; font-size: 12px;")
        side_layout.addWidget(self.lbl_status)

        self.graph_layout = pg.GraphicsLayoutWidget()
        self.graph_layout.setBackground('#1e1e1e')
        
        main_layout.addWidget(side_panel)
        main_layout.addWidget(self.graph_layout, 1)

    def generate_test_csv(self):
        file_path, _ = QFileDialog.getSaveFileName(self, "保存测试数据", "test_grid.csv", "CSV Files (*.csv)")
        if not file_path: return

        self.lbl_status.setText("生成数据中...")
        QApplication.processEvents()
        try:
            rows = 100000
            t = np.linspace(0, 10, rows)
            df = pd.DataFrame({
                'Time': t,
                'Vol_A': np.sin(2 * np.pi * 5 * t),
                'Vol_B': np.cos(2 * np.pi * 5 * t),
                'Amp': np.random.normal(0, 0.5, rows) + 2,
                'Temp': t * 2 + np.random.rand(rows)
            })
            df.to_csv(file_path, index=False)
            self.lbl_status.setText("生成完毕")
            QMessageBox.information(self, "完成", "测试文件已生成。")
        except Exception as e:
            QMessageBox.critical(self, "错误", str(e))

    def load_csv(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "打开 CSV", "", "CSV (*.csv)")
        if not file_path: return

        self.lbl_status.setText(f"读取中: {os.path.basename(file_path)}...")
        self.setCursor(Qt.CursorShape.WaitCursor)
        QApplication.processEvents()

        try:
            self.df = pd.read_csv(file_path)
            self.column_state_cache = {} 
            self.lbl_status.setText(f"已加载: {len(self.df)} 行")
            self.btn_select_cols.setEnabled(True)
            self.open_column_selector()
        except Exception as e:
            self.lbl_status.setText("加载失败")
            QMessageBox.critical(self, "错误", f"无法读取文件:\n{e}")
        finally:
            self.setCursor(Qt.CursorShape.ArrowCursor)

    def open_column_selector(self):
        if self.df is None: return
        dialog = ColumnSelectorDialog(list(self.df.columns), self.column_state_cache, self)
        if dialog.exec():
            self.column_state_cache = dialog.new_cache_config
            self.plot_data(dialog.selected_x, dialog.grouped_selection)

    def plot_data(self, x_col, grouped_selection):
        self.graph_layout.clear()
        self.plot_items = []
        self.lbl_status.setText("正在渲染...")
        QApplication.processEvents()

        try:
            x_data = self.df[x_col].to_numpy()
            sorted_row_ids = sorted(grouped_selection.keys())
            
            for i, row_id in enumerate(sorted_row_ids):
                items_in_this_row = grouped_selection[row_id]
                
                p = self.graph_layout.addPlot(row=i, col=0)
                p.setClipToView(True)
                
                # --- 核心修改：始终开启网格 ---
                p.showGrid(x=True, y=True, alpha=0.3)
                
                p.addLegend(offset=(10, 10))

                col_names = [item['name'] for item in items_in_this_row]
                y_label = " & ".join(col_names)
                if len(y_label) > 30: y_label = f"Group {row_id}"
                p.setLabel('left', y_label)

                # --- 核心修改：处理底部轴的显示 ---
                if i == len(sorted_row_ids) - 1:
                    # 如果是最底下的图，显示完整的轴（带数字标签）
                    p.setLabel('bottom', x_col)
                else:
                    # 如果是上面的图：
                    # 1. 必须 showAxis('bottom')，否则网格线没有刻度依据，画不出来
                    p.showAxis('bottom', True)
                    # 2. 隐藏数字 (showValues=False)，保持界面整洁
                    p.getAxis('bottom').setStyle(showValues=False)
                    # 3. 清空文字标签
                    p.setLabel('bottom', '')
                    
                    # 设置联动
                    if self.plot_items:
                        p.setXLink(self.plot_items[0])

                for item in items_in_this_row:
                    col_name = item['name']
                    hex_color = item['color']
                    y_data = self.df[col_name].to_numpy()
                    
                    pen = pg.mkPen(color=hex_color, width=1.5)
                    curve = p.plot(x_data, y_data, pen=pen, name=col_name)
                    curve.setDownsampling(ds=True, auto=True, method='peak')

                self.plot_items.append(p)

            if len(self.plot_items) > 1:
                for p in self.plot_items[1:]:
                    p.setXLink(self.plot_items[0])

            self.lbl_status.setText(f"渲染完成: {len(sorted_row_ids)} 行图表")

        except Exception as e:
            import traceback
            traceback.print_exc()
            QMessageBox.critical(self, "绘图错误", str(e))
            self.lbl_status.setText("绘图出错")

    def reset_views(self):
        if self.plot_items:
            self.plot_items[0].enableAutoRange()

if __name__ == "__main__":
    app = QApplication(sys.argv)
    if hasattr(Qt.ApplicationAttribute, 'AA_UseHighDpiPixmaps'):
        app.setAttribute(Qt.ApplicationAttribute.AA_UseHighDpiPixmaps, True)
    
    window = ProScopeFullGrid()
    window.show()
    sys.exit(app.exec())